# -*- coding: utf-8 -*-
"""
Created on Tue Nov 14 10:19:43 2017

@author: Luther
"""

import tushare as ts
import pandas as pd
import numpy as np
from sklearn import linear_model
import matplotlib.pyplot as plt


#获取基于上证指数测算的证券市场收益率
def get_szzs():
    sh = ts.get_k_data(
        '000001', index=True, start='2016-01-01', end='2016-12-31')
    sh.index = sh.date
    sh = sh.drop(['date'], axis=1)
    sh['ln_i_close'] = np.log(sh.close)
    x = pd.DataFrame(sh.ln_i_close.diff()[1:])
    return x


#计算个股beta系数
def get_beta(code, x):

    stock = ts.get_k_data(str(code), start='2016-01-01', end='2016-12-31')
    stock.index = stock.date
    stock = stock.drop(['date'], axis=1)
    stock['ln_close'] = np.log(stock.close)
    y = pd.DataFrame(stock.ln_close.diff()[1:])
    reg = pd.DataFrame.join(y, x)
    reg.columns = ['y', 'x']

    model = linear_model.LinearRegression()
    model.fit(reg.x.values.reshape(-1, 1), reg.y)
    beta = model.coef_[0]

    return beta


#获取上证50指数成分股公司的beta系数
sz50 = ts.get_sz50s()
sz50.index = sz50.code
sz50 = pd.DataFrame(sz50.name)

industry = ts.get_industry_classified()
industry.index = industry.code
industry = pd.DataFrame(industry.c_name)

result = pd.DataFrame.join(sz50, industry)
result = result.dropna()

x = get_szzs()
beta_list = []
for i in range(len(result)):
    try:
        beta = get_beta(result.index[i], x)
    except:
        beta = 0
    beta_list.append(beta)
result['beta'] = beta_list
result = result[result.beta != 0]

#统计上证50指数成分股在不同行业的分布
count = result.groupby('c_name').name.count()
countdf = pd.DataFrame()
countdf['行业'] = count.index
countdf['上证50公司数量'] = count.values
countdf = countdf.sort_values(by='上证50公司数量', ascending=False)
countdf.index = np.arange(1, len(count) + 1)
print(countdf)

plt.barh(np.arange(len(count)), countdf['上证50公司数量'])
plt.yticks(np.arange(len(count)), countdf['行业'], fontproperties='SimHei')
plt.ylabel('行业', fontproperties='SimHei')
plt.xlabel('上证50指数成分股数量', fontproperties='SimHei')
for i, j in zip(countdf['上证50公司数量'], range(len(countdf))):
    plt.text(i + 0.5, j, i, ha='center', va='center')
plt.title('上证50指数成分股行业分布图', fontproperties='SimHei')
plt.show()

#统计上证50指数成分股分行业beta系数
beta_mean = result.groupby('c_name').beta.mean()
betadf = pd.DataFrame()
betadf['行业'] = beta_mean.index
betadf['β系数'] = beta_mean.values
betadf = betadf.sort_values(by='β系数', ascending=False)
betadf.index = np.arange(1, len(beta_mean) + 1)
print('\n')
print(betadf)

plt.barh(np.arange(len(beta_mean)), betadf['β系数'])
plt.yticks(np.arange(len(beta_mean)), beta_mean.index, fontproperties='SimHei')
plt.ylabel('行业', fontproperties='SimHei')
plt.xlabel('β系数', fontproperties='SimHei')
plt.xlim(xmax=1.6)
for i, j in zip(betadf['β系数'], range(len(betadf))):
    plt.text(i + 0.06, j, '{:.2f}'.format(i), ha='center', va='center')
plt.title('上证50指数成分股行业β系数图', fontproperties='SimHei')
plt.show()
